import datetime
import functools
import logging

import altair as alt
from absl import app

import experimental.prophet as prophet
from coin.exchange.bitmex.kr_rest.futures_product import BitmexFuturesProduct
from coin.exchange.okex.kr_rest.product import OkexProduct
from coin.exchange.upbit_v1.kr_rest.product import UpbitProduct
from experimental.prophet.ops import timeseries


def model(coin):
  ts = prophet.global_timestamp()
  timer_1s = prophet.timer('1s')

  product = BitmexFuturesProduct.FromStr('%s-USD.PERPETUAL' % coin)
  feed = prophet.fastfeed_coin(product, 'Futures.Bitmex')
  book = prophet.get_book_from_fastfeed(feed)
  instrument = prophet.get_instrument_from_fastfeed(feed)

  midp = (book.ask0.price + book.bid0.price) / 2.

  past_return_5s = midp - timeseries.time_shift(midp, '5s')
  past_return_10s = midp - timeseries.time_shift(midp, '10s')
  past_return_30s = midp - timeseries.time_shift(midp, '30s')

  return_5s = prophet.time_shift_y(past_return_5s, '5s').named_ref('return_5s')
  return_10s = prophet.time_shift_y(past_return_10s, '10s').named_ref('return_10s')
  return_30s = prophet.time_shift_y(past_return_30s, '30s').named_ref('return_30s')

  open_interest_diff_5s = (
      instrument.open_interest
      - timeseries.time_shift(instrument.open_interest, '5s')).named_ref('open_interest_5s')
  open_interest_diff_10s = (
      instrument.open_interest
      - timeseries.time_shift(instrument.open_interest, '10s')).named_ref('open_interest_10s')
  open_interest_diff_30s = (
      instrument.open_interest
      - timeseries.time_shift(instrument.open_interest, '30s')).named_ref('open_interest_30s')

  with prophet.control_if(ts % 10**9 == 0):
    aggregator = prophet.scalar_value_aggregator([
        ts,
        return_5s,
        return_10s,
        return_30s,
        open_interest_diff_5s,
        open_interest_diff_10s,
        open_interest_diff_30s
    ])

  return aggregator


def main(argv):
  # coins = ['BTC', 'ETH', 'XRP', 'BCHABC', 'LTC', 'EOS']
  coin = 'BTC'
  df = prophet.run_from_fastfeed_multiprocess(functools.partial(model, coin),
                                              datetime.date(2019, 7, 1),
                                              datetime.datetime(2019, 7, 1, 2),
                                              machine='feed-01.ap-northeast-1.aws')

  n_col = len(df.columns)
  charts = []
  for i in range(1, 4):
    for j in range(4, n_col):
      chart = alt.Chart(df).mark_circle(size=2).encode(x=df.columns[i], y=df.columns[j])
      charts.append(chart)

  alt.concat(*charts, columns=2).save('chart.html')

  print(df.head())


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG')
  app.run(main)
